from math import gcd
from random import randint
from utils import generatePrime, noneHash, invmod, is2PrimitiveRoot

class ElGamal:
    def __init__(self) -> None:
        pass

    def generateKeys(self, size: int) -> tuple:
        """
        随机生成ElGamal算法的公钥对和私钥对
        输入 大素数p的位数
        返回 (y, p, g) 和 x
        """
        # 生成一个最小本原根是2的大素数
        p = generatePrime(size)
        while not is2PrimitiveRoot(p):
            p = generatePrime(size)
            
        g = 2
        x = randint(2, p-1)
        y = pow(g, x, p)
        return (y, p, g), x    

    def signature(self, keys: tuple, m: int, hash=noneHash) -> tuple:
        """
        对消息m使用keys进行ElGamal数字签名
        输入 消息m 公私钥对keys 签名哈希函数hash
        返回 签名 (r,s)
        """
        y, p, g = keys[0]
        x = keys[1]
        self.keys = keys

        # 随机产生奇数k 
        # p-1一定是偶数 所以产生的k一定要是奇数才有可能和p-1互质
        k = randint(1, p-1) | 1
        while gcd(k, p-1) != 1:
            k = randint(1, p-1) | 1
        self.k = k

        r = pow(g, k, p)
        s = (invmod(k, p-1) * (hash(m) - x*r)) % (p-1)

        return (r,s) 
            
    def verify(self, keys: tuple, m: int, sig: tuple , hash=noneHash) -> bool:
        """
        对消息m使用keys进行ElGamal数字签名认证
        输入 消息m 公私钥对keys 签名sig 签名哈希函数hash
        返回 是否是有效签名
        """
        self.keys = keys
        self.sig = sig
        y, p, g = keys[0]
        r, s = sig
        return ((pow(y,r) * pow(r,s)) % p) == (pow(g,hash(m)) % p)
    
    def showDetails(self) -> None:
        """
        将签名过程中的数据打印出来
        """
        print("公钥对:", self.keys[0])
        print("私钥:", self.keys[1])
        print("随机数k", self.k)
        print("签名结果", self.sig)
        print()

        
# for debug
# elg = ElGamal()
# keys, m = elg.generateKeys(10), 200110720
# sig = elg.signature(keys, m)
# print(elg.verify(keys, m, sig))
# elg.showDetails()